
use image::{DynamicImage, GenericImageView, GrayImage, ImageBuffer, Luma};
use num_integer::Roots;
/// 扫描件预处理主函数，针对扫描件特点优化
pub fn preprocess_scan_image(img: DynamicImage) -> Result<DynamicImage, Box<dyn std::error::Error>> {
    // 2. 转换为灰度图
    let gray_img = img.to_luma8();

    // 3. 去除扫描件边框（优化版本，避免裁剪文字）
    let cropped_img = remove_borders_improved(&gray_img, 30, 5)?;

    // // 4. 检测并校正倾斜
    // let rotated_img = correct_skew(&cropped_img)?;

    // // 5. 背景归一化（处理扫描件常见的深浅不均）
    // let normalized_img = normalize_background(&cropped_img, 20);

    // 6. 全局阈值二值化（扫描件通常对比度较好）
    let binary_img = global_threshold(&cropped_img, 180);

    // // 7. 去除扫描斑点噪点
    // let denoised_img = remove_scan_noise(&binary_img);

    Ok(DynamicImage::ImageLuma8(binary_img))
}

/// 改进的边框去除函数
/// sensitivity: 边框检测灵敏度（0-255），值越大越容易识别边框
/// min_border_width: 最小边框宽度，小于此宽度的不视为边框
fn remove_borders_improved(
    img: &GrayImage,
    sensitivity: u8,
    min_border_width: u32,
) -> Result<GrayImage, Box<dyn std::error::Error>> {
    let (width, height) = img.dimensions();

    // 首先进行边缘检测，突出边框和文字边缘
    let edges = sobel_edge_detection(img);

    // 计算边缘密度，边框区域通常边缘密度低
    let edge_density = calculate_edge_density(&edges, 5);

    // 检测上下左右四个边的边框
    let left_border = detect_side_border(&edge_density, true, true, sensitivity, min_border_width)?;
    let right_border =
        detect_side_border(&edge_density, true, false, sensitivity, min_border_width)?;
    let top_border = detect_side_border(&edge_density, false, true, sensitivity, min_border_width)?;
    let bottom_border =
        detect_side_border(&edge_density, false, false, sensitivity, min_border_width)?;

    // 计算裁剪区域，保留安全边距
    let margin = 3;
    let x1 = left_border + margin;
    let x2 = width - right_border - margin;
    let y1 = top_border + margin;
    let y2 = height - bottom_border - margin;

    // 确保裁剪区域有效
    if x1 >= x2 || y1 >= y2 {
        // 如果裁剪无效，返回原图
        return Ok(img.clone());
    }

    // 裁剪图像
    let cropped = img.view(x1, y1, x2 - x1, y2 - y1).to_image();
    Ok(cropped)
}

/// 计算边缘密度
fn calculate_edge_density(edges: &GrayImage, window_size: u32) -> GrayImage {
    let (width, height) = edges.dimensions();
    let mut density = ImageBuffer::new(width, height);
    let half = window_size / 2;

    for y in 0..height {
        for x in 0..width {
            let mut edge_count = 0u32;
            let mut total = 0u32;

            // 计算窗口内的边缘密度
            for dy in 0..window_size {
                for dx in 0..window_size {
                    let nx = x as i32 + dx as i32 - half as i32;
                    let ny = y as i32 + dy as i32 - half as i32;

                    if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
                        if edges.get_pixel(nx as u32, ny as u32)[0] > 100 {
                            edge_count += 1;
                        }
                        total += 1;
                    }
                }
            }

            // 边缘密度转换为0-255值
            let density_value = if total == 0 {
                0
            } else {
                (255.0 * (edge_count as f64 / total as f64)) as u8
            };

            density.put_pixel(x, y, Luma([density_value]));
        }
    }

    density
}

/// 检测单边边框
/// is_vertical: 是否为垂直边框（左右）
/// is_start: 是否为起始边（左/上）
fn detect_side_border(
    edge_density: &GrayImage,
    is_vertical: bool,
    is_start: bool,
    sensitivity: u8,
    min_border_width: u32,
) -> Result<u32, Box<dyn std::error::Error>> {
    let (width, height) = edge_density.dimensions();
    let max_length = if is_vertical { width } else { height };
    let step = if is_start { 1 } else { max_length - 1 };
    let end = if is_start { max_length } else { 0 };
    let direction = if is_start { 1 } else { -1 };

    // 逐行/逐列检测边框
    let mut border_width = 0;
    let mut consecutive_matches = 0;

    let mut pos = step as i32;
    while if is_start {
        pos < end as i32
    } else {
        pos >= end as i32
    } {
        let p = pos as u32;
        let mut is_border = true;
        let mut count = 0;
        let mut border_count = 0;

        // 检查整行/整列是否为边框
        for i in 0..if is_vertical { height } else { width } {
            let (x, y) = if is_vertical { (p, i) } else { (i, p) };
            let density = edge_density.get_pixel(x, y)[0];

            // 边框区域边缘密度低
            if density < sensitivity {
                border_count += 1;
            }
            count += 1;
        }

        // 如果超过70%的区域符合边框特征，则认为是边框
        if border_count as f64 / count as f64 > 0.7 {
            consecutive_matches += 1;
        } else {
            // 如果连续不匹配次数超过3，则认为边框结束
            if consecutive_matches > 0 {
                consecutive_matches = 0;
                break;
            }
        }

        border_width += 1;
        pos += direction;
    }

    // 确保边框宽度达到最小要求，否则不视为边框
    if border_width < min_border_width {
        border_width = 0;
    }

    Ok(border_width)
}

// 以下函数与之前版本相同，为保持完整性列出
fn correct_skew(img: &GrayImage) -> Result<GrayImage, Box<dyn std::error::Error>> {
    let (width, height) = img.dimensions();

    // 边缘检测（使用简单的Sobel算子）
    let edges = sobel_edge_detection(img);

    // 霍夫变换检测线条，找出主要倾斜角度
    let angle = hough_transform_detect_angle(&edges, -15.0, 15.0, 0.5)?;

    // 如果倾斜角度很小，不需要旋转
    if angle.abs() < 0.5 {
        return Ok(img.clone());
    }

    // 旋转图像校正倾斜
    let rotated = rotate_image(img, angle);
    Ok(rotated)
}

fn sobel_edge_detection(img: &GrayImage) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut edges = ImageBuffer::new(width, height);

    // Sobel算子卷积核
    let gx = [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]];
    let gy = [[-1, -2, -1], [0, 0, 0], [1, 2, 1]];

    for y in 1..height - 1 {
        for x in 1..width - 1 {
            let mut sum_x = 0i32;
            let mut sum_y = 0i32;

            // 应用卷积
            for dy in -1..=1 {
                for dx in -1..=1 {
                    let pixel =
                        img.get_pixel((x as i32 + dx) as u32, (y as i32 + dy) as u32)[0] as i32;
                    sum_x += pixel * gx[(dy + 1) as usize][(dx + 1) as usize];
                    sum_y += pixel * gy[(dy + 1) as usize][(dx + 1) as usize];
                }
            }

            // 计算梯度幅度
            let magnitude = (sum_x.pow(2) + sum_y.pow(2)).sqrt() as u8;
            edges.put_pixel(x, y, Luma([magnitude]));
        }
    }

    edges
}

fn hough_transform_detect_angle(
    edges: &GrayImage,
    min_angle: f64,
    max_angle: f64,
    angle_step: f64,
) -> Result<f64, Box<dyn std::error::Error>> {
    let (width, height) = edges.dimensions();
    let center_x = width as f64 / 2.0;
    let center_y = height as f64 / 2.0;

    // 收集边缘点
    let mut edge_points = Vec::new();
    for y in 0..height {
        for x in 0..width {
            if edges.get_pixel(x, y)[0] > 100 {
                // 边缘点阈值
                edge_points.push((x as f64 - center_x, y as f64 - center_y));
            }
        }
    }

    if edge_points.is_empty() {
        return Err("未检测到边缘点，无法计算倾斜角度".into());
    }

    // 计算每个角度的投票
    let mut max_votes = 0;
    let mut best_angle = 0.0;

    let num_steps = ((max_angle - min_angle) / angle_step) as usize;
    for i in 0..=num_steps {
        let angle = min_angle + (i as f64) * angle_step;
        let radians = angle.to_radians();

        // 计算该角度下的线条投票
        let mut votes = 0;
        for &(x, y) in &edge_points {
            // 点到直线的距离公式（简化版）
            let distance = (y * radians.cos() - x * radians.sin()).abs();
            if distance < 2.0 {
                // 距离阈值，认为点在直线上
                votes += 1;
            }
        }

        if votes > max_votes {
            max_votes = votes;
            best_angle = angle;
        }
    }

    Ok(best_angle)
}

fn rotate_image(img: &GrayImage, angle: f64) -> GrayImage {
    let (width, height) = img.dimensions();
    let radians = angle.to_radians();

    // 计算旋转后的图像尺寸
    let cos_theta = radians.cos().abs();
    let sin_theta = radians.sin().abs();
    let new_width = (width as f64 * cos_theta + height as f64 * sin_theta) as u32;
    let new_height = (width as f64 * sin_theta + height as f64 * cos_theta) as u32;

    let mut rotated = ImageBuffer::new(new_width, new_height);
    let center_x = width as f64 / 2.0;
    let center_y = height as f64 / 2.0;
    let new_center_x = new_width as f64 / 2.0;
    let new_center_y = new_height as f64 / 2.0;

    // 反向映射：对新图像的每个像素，找到原始图像中的对应位置
    for y in 0..new_height {
        for x in 0..new_width {
            // 平移到原点
            let x_prime = x as f64 - new_center_x;
            let y_prime = y as f64 - new_center_y;

            // 旋转（逆变换）
            let x_orig = x_prime * radians.cos() - y_prime * (-radians.sin()) + center_x;
            let y_orig = x_prime * (-radians.sin()) + y_prime * radians.cos() + center_y;

            // 双线性插值获取像素值
            let pixel = if x_orig >= 0.0
                && x_orig < width as f64 - 1.0
                && y_orig >= 0.0
                && y_orig < height as f64 - 1.0
            {
                bilinear_interpolation(img, x_orig, y_orig)
            } else {
                255 // 边缘填充白色
            };

            rotated.put_pixel(x, y, Luma([pixel]));
        }
    }

    rotated
}

fn bilinear_interpolation(img: &GrayImage, x: f64, y: f64) -> u8 {
    let x0 = x.floor() as u32;
    let x1 = x0 + 1;
    let y0 = y.floor() as u32;
    let y1 = y0 + 1;

    let dx = x - x0 as f64;
    let dy = y - y0 as f64;

    let p00 = img.get_pixel(x0, y0)[0] as f64;
    let p10 = img.get_pixel(x1, y0)[0] as f64;
    let p01 = img.get_pixel(x0, y1)[0] as f64;
    let p11 = img.get_pixel(x1, y1)[0] as f64;

    // 双线性插值计算
    let p0 = p00 * (1.0 - dx) + p10 * dx;
    let p1 = p01 * (1.0 - dx) + p11 * dx;
    let p = p0 * (1.0 - dy) + p1 * dy;

    p as u8
}

fn normalize_background(img: &GrayImage, window_size: u32) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut result = ImageBuffer::new(width, height);
    let half_window = (window_size / 2) as i32;

    for y in 0..height {
        for x in 0..width {
            // 计算局部窗口的平均值（背景估计）
            let mut sum = 0u32;
            let mut count = 0u32;

            for dy in -half_window..=half_window {
                for dx in -half_window..=half_window {
                    let nx = x as i32 + dx;
                    let ny = y as i32 + dy;

                    if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
                        sum += img.get_pixel(nx as u32, ny as u32)[0] as u32;
                        count += 1;
                    }
                }
            }

            if count == 0 {
                result.put_pixel(x, y, Luma([255]));
                continue;
            }

            let mean = sum / count;
            let pixel = img.get_pixel(x, y)[0] as u32;

            // 背景归一化公式：新值 = 255 * (原值 / 背景估计值)，限制在0-255
            let normalized = (255.0 * (pixel as f64 / mean as f64)).clamp(0.0, 255.0) as u8;
            result.put_pixel(x, y, Luma([normalized]));
        }
    }

    result
}

fn global_threshold(img: &GrayImage, threshold: u8) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut result = ImageBuffer::new(width, height);

    for (x, y, pixel) in img.enumerate_pixels() {
        let value = if pixel[0] < threshold { 0 } else { 255 };
        result.put_pixel(x, y, Luma([value]));
    }

    result
}

fn remove_scan_noise(img: &GrayImage) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut result = img.clone();

    // 首先进行腐蚀操作，去除小的白色噪点
    let eroded = erode(img, 1);

    // 然后进行膨胀操作，恢复文字形状
    dilate(&eroded, 1)
}

fn erode(img: &GrayImage, kernel_size: u32) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut result = ImageBuffer::new(width, height);
    let k = kernel_size as i32;

    for y in 0..height {
        for x in 0..width {
            let mut has_black = false;

            // 检查核内是否有黑色像素
            for dy in -k..=k {
                for dx in -k..=k {
                    let nx = x as i32 + dx;
                    let ny = y as i32 + dy;

                    if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
                        if img.get_pixel(nx as u32, ny as u32)[0] == 0 {
                            has_black = true;
                            break;
                        }
                    }
                }
                if has_black {
                    break;
                }
            }

            // 腐蚀规则：如果核内有黑色，则保留黑色，否则变为白色
            let value = if has_black { 0 } else { 255 };
            result.put_pixel(x, y, Luma([value]));
        }
    }

    result
}

fn dilate(img: &GrayImage, kernel_size: u32) -> GrayImage {
    let (width, height) = img.dimensions();
    let mut result = ImageBuffer::new(width, height);
    let k = kernel_size as i32;

    for y in 0..height {
        for x in 0..width {
            let mut has_white = false;

            // 检查核内是否有白色像素
            for dy in -k..=k {
                for dx in -k..=k {
                    let nx = x as i32 + dx;
                    let ny = y as i32 + dy;

                    if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
                        if img.get_pixel(nx as u32, ny as u32)[0] == 255 {
                            has_white = true;
                            break;
                        }
                    }
                }
                if has_white {
                    break;
                }
            }

            // 膨胀规则：如果核内有白色，则变为白色，否则保留黑色
            let value = if has_white { 255 } else { 0 };
            result.put_pixel(x, y, Luma([value]));
        }
    }

    result
}
